import os
import logging
import pandas as pd
import torch
import time
import data_augment
from transformers import BertTokenizer
from bert import Bert
from train_eval import train
from utils import config_to_json_string, read_csv


# 定义RoBERTa模型配置类，用于存储模型训练和运行的相关参数
class RobertaPairConfig:

    def __init__(self):
        # 预训练模型相关路径配置
        self.pretrain_path = '../pretrain_models/roberta_large_pair'
        _config_file = 'bert_config.json'  # 模型配置文件
        _model_file = 'pytorch_model.bin'  # 预训练模型权重文件
        _tokenizer_file = 'vocab.txt'  # 分词器词汇表文件
        self.config_file = os.path.join(self.pretrain_path, _config_file)
        self.model_name_or_path = os.path.join(self.pretrain_path, _model_file)
        self.tokenizer_file = os.path.join(self.pretrain_path, _tokenizer_file)

        # 数据路径配置
        self.train_path = '../data/KUAKE/KUAKE-QQR_train.json'  # 训练数据路径
        self.dev_path = '../data/KUAKE/KUAKE-QQR_dev.json'  # 验证数据路径
        self.test_path = '../data/KUAKE/KUAKE-QQR_test.json'  # 测试数据路径
        self.aug_data_path = '../data/KUAKE/KUAKE-QQR_augment.csv'  # 增强数据路径

        # 模型配置
        self.use_model = 'bert'  # 使用的模型类型
        self.models_name = 'roberta_large_pair'  # 模型名称
        self.task = 'KUAKE'  # 任务名称
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 运行设备
        self.requires_grad = True  # 是否需要计算梯度
        self.class_list = ['0', '1', '2']  # 分类类别列表
        self.num_labels = len(self.class_list)  # 分类类别数量
        self.train_num_examples = 0  # 训练样本数量
        self.dev_num_examples = 0  # 验证样本数量
        self.test_num_examples = 0  # 测试样本数量
        self.hidden_dropout_prob = 0.1  # 隐藏层dropout概率
        self.multi_drop = 5  # 多次dropout的次数
        self.hidden_size = 1024  # 隐藏层大小
        self.early_stop = False  # 是否启用早停
        self.require_improvement = 800  # 早停判断的步数
        self.num_train_epochs = 5  # 训练轮数
        self.batch_size = 64  # 批次大小
        self.pad_size = 64  # 输入序列的填充长度
        self.learning_rate = 8e-6  # 学习率
        self.head_learning_rate = 1e-3  # 分类层的学习率
        self.weight_decay = 0.01  # 权重衰减系数
        self.warmup_proportion = 0.1  # 学习率预热比例

        # 日志配置
        self.is_logging2file = True  # 是否将日志写入文件
        self.logging_dir = '../logging' + '/' + self.models_name  # 日志目录

        # 模型保存配置
        self.save_path = '../my_model'  # 模型保存路径
        self.save_file = self.models_name  # 模型保存文件名

        # 数据增强配置
        self.data_augment = True  # 是否使用数据增强

        # 差分学习率配置
        self.diff_learning_rate = False  # 是否使用差分学习率

        # 预处理配置
        self.stop_word_valid = True  # 是否使用停用词过滤

        # 新增属性：是否显示进度条
        self.show_progress = True


# 主程序入口
if __name__ == '__main__':
    # 初始化配置
    config = RobertaPairConfig()
    print(f"运行设备: {config.device}")

    # 配置日志
    logging_filename = None
    if config.is_logging2file is True:
        # 创建以当前时间命名的日志文件
        file = time.strftime('%Y-%m-%d_%H-%M-%S') + '.log'
        logging_filename = os.path.join(config.logging_dir, file)
        # 创建日志目录(如果不存在)
        if not os.path.exists(config.logging_dir):
            os.makedirs(config.logging_dir)
    # 配置日志格式和级别
    logging.basicConfig(filename=logging_filename, format='%(levelname)s: %(message)s', level=logging.INFO)

    # 初始化分词器
    tokenizer = BertTokenizer.from_pretrained(config.tokenizer_file)

    # 读取训练数据
    train_examples = read_csv(config.train_path)

    # 如果启用数据增强，则加载增强数据并添加到训练数据中
    if config.data_augment:
        # 从CSV文件读取增强数据并转换为列表格式
        augment_examples = data_augment.dataframe_to_list(pd.read_csv(config.aug_data_path))
        train_examples.extend(augment_examples)

    # 读取验证数据
    dev_examples = read_csv(config.dev_path)

    # 初始化模型
    model = Bert(config)

    # 记录配置信息到日志
    logging.info("模型配置: %s", config_to_json_string(config))

    # 开始训练模型
    train(
        config=config,  # 模型配置
        model=model,  # 模型实例
        tokenizer=tokenizer,  # 分词器
        train_data=train_examples,  # 训练数据
        dev_data=dev_examples,  # 验证数据
    )